# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from typing import Optional
from typing import Tuple

from ppsci.equation import PDE


class Convection(PDE):
    r"""Class for 1D convection equation.

    $$
    u_t - c \nabla u = 0
    $$

    Args:
        dim (int): Dimension of the equation. Defaults to 1.
        time (bool): Whether to include time dimension. Defaults to True.
        c (float): Convection coefficient. Defaults to 1.0.
        detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing.
            Defaults to None.

    Examples:
        >>> from pdeapp.equation.pde import Convection1D
        >>> pde = Convection1D(dim=1, time=True, c=0.1)
        >>> print(pde.equations)
    """

    def __init__(
        self,
        dim: int = 1,
        time: bool = True,
        c: float = 1.0,
        detach_keys: Optional[Tuple[str, ...]] = None,
    ):
        super().__init__()
        self.c = c
        self.dim = dim
        self.time = time

        t, x, y, z = self.create_symbols("t x y z")
        if dim == 1:
            invars = (x,)
        elif dim == 2:
            invars = (x, y)
        elif dim == 3:
            invars = (x, y, z)
        else:
            raise ValueError("Dimension should be 1, 2 or 3.")

        if time:
            invars = (t,) + invars

        u = self.create_function("u", invars)

        if time:
            u_t = u.diff(t)
            coord_var = invars[1:]
        else:
            u_t = 0
            coord_var = invars

        conv = 0
        for invar in coord_var:
            conv += u.diff(invar)

        self.add_equation("convection", u_t + c * conv)

        self._apply_detach()
